#include <algorithm/analytical_model/algorithm/cuda/Utils.h>
#include <algorithm/analytical_model/algorithm/cuda/evaluation/metrics/RAD.cuh>

#include <cuda_runtime.h>
#include <thrust/async/for_each.h>
#include <thrust/async/reduce.h>
#include <thrust/async/transform.h>
#include <thrust/device_ptr.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/tuple.h>
#include <vector_functions.h>

#include <cmath>

using namespace solar::cuda;

auto solar::cuda::metrics::RAD(float* d_input, float* d_ground_truth,
                               const inter_op::Int2& resolution, const cudaStream_t stream) -> float
{
    auto result = metrics::async::RAD(d_input, d_ground_truth, resolution, stream);
    return std::get<0>(result).get() / std::get<1>(result).get();
}

auto solar::cuda::metrics::async::RAD(float* d_input, float* d_ground_truth,
                                      const inter_op::Int2& resolution, const cudaStream_t stream)
    -> std::tuple<thrust::cuda::unique_eager_future<float>,
                  thrust::cuda::unique_eager_future<float>>
{
    const std::size_t total_pixel_num = resolution.x * resolution.y;
    const thrust::device_ptr<float> ground_truth_ptr(d_ground_truth);
    auto zip_iterator = thrust::make_zip_iterator(d_input, d_ground_truth);
    auto ground_truth_sum = thrust::async::reduce(
        thrust::device.on(stream), d_ground_truth, d_ground_truth + total_pixel_num, 0.0F,
        [] __host__ __device__(const float& num1, const float& num2)
        { return (num1 + num2); });
    float* result_diff = nullptr;
    CHECK_CUDA_ERRORS(
        cudaMallocAsync(std::addressof(result_diff), total_pixel_num * sizeof(float), stream));
    thrust::async::transform(thrust::device.on(stream), zip_iterator,
                             zip_iterator + total_pixel_num, result_diff,
                             [] __host__ __device__(const thrust::tuple<float, float>& iter)
                             { return std::abs(thrust::get<0>(iter) - thrust::get<1>(iter)); });
    auto diff_sum =
        thrust::async::reduce(thrust::device.on(stream), result_diff, result_diff + total_pixel_num,
                              0.0F, [] __host__ __device__(const float& num1, const float& num2)
                              { return (num1 + num2); });
    CHECK_CUDA_ERRORS(cudaFreeAsync(result_diff, stream));
    return {std::move(diff_sum), std::move(ground_truth_sum)};
}
